import xarray as xr
import numpy as np
from ..physics.unit import _handle_quantity, unit, Q_
from ..physics.constant import EARTH_GRAVITY, ROSSBY_PARAMETER_ON_EQUATOR, EARTH_EQUATOR_LENGTH
from pint import Quantity
from shapely.geometry import Polygon, Point
from shapely import points as spoints
from typing import Literal
import json

PI2 = np.pi * 2

def _121_smooth(data, times=10):
    from scipy.signal import convolve2d
    kernel = np.array([[0, 1, 0], [1, 2, 1], [0, 1, 0]]) / 6
    for i in range(times):
        data = convolve2d(data, kernel, mode='same')
    return data


def wk99_spectrum(data: xr.DataArray,
                  lon_name: str = "longitude",
                  lat_name: str = "latitude",
                  time_name: str = "time",
                  cycle_days: int = 96,
                  cycle_overlap_days: int = 64,
                  bg_smooth_times: int = 10,) -> xr.DataArray:
    """
    计算变量的频率-波数谱，计算方法按照 WK99。需要输入一个三维的DataArray，
    返回一个频率-波数谱的 DataArray


    WHEELER M, KILADIS G N, 1999. Convectively Coupled Equatorial Waves:
    Analysis of Clouds and Temperature in the Wavenumber–Frequency Domain[J/OL].
    Journal of the Atmospheric Sciences, 56(3): 374-399.

    Parameters
    ----------
    data : xr.DataArray
        三维的时间、经纬、纬度数据

        3D data with time, longitude and latitude dimensions
    lon_name : str, optional
        data 中经度的名称  The name of longitude in data
    lat_name : str, optional
        data 中纬度的名称  The name of latitude in data
    time_name : str, optional
        data 中时间的名称  The name of time in data
    cycle_days : int, optional
        循环计算 FFT 的时间，默认为 96 日

        The time for FFT calculation, default is 96 days
    cycle_overlap_days : int, optional
        循环重叠的天数，默认为 64 日

        The time for FFT calculation, default is 64 days
    bg_smooth_times : int, optional
        背景平滑的次数，默认为 10 次

        The number of times to smooth the background, default is 10 times
    Returns
    -------
    xr.DataArray
        返回一个三维数据：(type, freq, wavenumber)，
        type: 对称和非对称的分量，分别为 'sym' 和 'asym'
        freq: 频率，单位为 1/day
        wavenumber: 波数，单位为赤道上的波数
    """
    from scipy.signal import windows, detrend

    # 转为 时间、纬度、经度
    data = data.transpose(time_name, lat_name, lon_name)

    dt = data[time_name].values[1] - data[time_name].values[0]
    dt = np.timedelta64(dt, 's')

    ticks_in_day = np.timedelta64(1, 'D') / dt

    cycle_length = int(cycle_days * ticks_in_day)
    cycle_overlap_length = int(cycle_overlap_days * ticks_in_day)

    data_north = data.sel({lat_name: slice(None, 0)}).values
    data_south = data.sel({lat_name: slice(0, None)}).values
    data_south = np.flip(data_south, axis=data.get_axis_num(lat_name))

    data_asym = (data_north - data_south) / 2
    data_sym = (data_north + data_south) / 2

    dt = int(cycle_length - cycle_overlap_length / 2)
    window = windows.tukey(cycle_length, 0.1)

    time_length = data_asym.shape[0]
    lon_length = data_asym.shape[-1]

    result = np.zeros((2, cycle_length, lon_length))

    for t in range(len(data_asym[::dt])):
        t_range = (t * dt, t * dt + cycle_length)
        if t_range[1] > time_length:
            break

        for i, d in enumerate([data_asym, data_sym]):
            d = d[t_range[0]:t_range[1]]
            d = d - np.mean(d, axis=0)
            d = detrend(d, axis=0)
            d = d * window[:, None, None]
            fft_result = np.fft.fft2(d, axes=(0, 2))
            fft_result = np.fft.fftshift(fft_result, axes=(0, 2))
            fft_result = fft_result / fft_result.shape[0] / fft_result.shape[2] * 2
            fft_result = np.mean(fft_result, axis=1)
            result[i] += np.abs(fft_result)**2

    result /= t

    # 计算波数
    lon_max = data[lon_name].values.max()
    lon_min = data[lon_name].values.min()
    lon_res = data[lon_name].values[1] - data[lon_name].values[0]
    lon_range = lon_max - lon_min

    zonal_res = data[lon_name].values[1] - data[lon_name].values[0]
    k = np.fft.fftfreq(data_asym.shape[-1], np.abs(zonal_res)/360)
    k = np.fft.fftshift(k)
    omega = np.fft.fftfreq(cycle_length, d=1/ticks_in_day)
    omega = np.fft.fftshift(omega) * -1

    background = np.abs(np.mean(result[:2], axis=0))
    background = _121_smooth(background, bg_smooth_times)
    result /= background

    omega_idx = np.where(omega < 0)[0][0]
    omega = omega[:omega_idx]
    result = result[:, :omega_idx, :]

    result = xr.DataArray(result, dims=["type", "freq", "wavenumber"],
                          coords={"type": ["asym", "sym", ],
                                  "freq": omega,
                                  "wavenumber": k},
                          attrs={
                              "wavenumber_units": "wave number in equator",
                              "freq_units": "1/day",
                          })

    return result


class FilterZone:
    def __init__(self, min_depth: Quantity | float, max_depth: Quantity | float,
                 min_freq: Quantity | float, max_freq: Quantity | float,
                 min_wavenumber: Quantity | float, max_wavenumber: Quantity | float,
                 g: Quantity | float = None, beta: Quantity | float = None):
        min_depth = _handle_quantity(min_depth, "m")
        max_depth = _handle_quantity(max_depth, "m")
        min_freq = _handle_quantity(min_freq, "1/day")
        max_freq = _handle_quantity(max_freq, "1/day")
        min_wavenumber = _handle_quantity(min_wavenumber, "zonal_wavenumber")
        max_wavenumber = _handle_quantity(max_wavenumber, "zonal_wavenumber")

        if g is not None:
            g = _handle_quantity(g, "m/s^2")
        else:
            g = EARTH_GRAVITY
        if beta is not None:
            beta = _handle_quantity(beta, "1/m/s")
        else:
            beta = ROSSBY_PARAMETER_ON_EQUATOR

        if min_depth >= max_depth:
            raise ValueError("min_depth should be less than max_depth")
        if min_freq >= max_freq:
            raise ValueError("min_freq should be less than max_freq")
        if min_wavenumber >= max_wavenumber:
            raise ValueError("min_wavenumber should be less than max_wavenumber")

        # 转换单位
        min_wavenumber = min_wavenumber.to("1/m") * 2 * np.pi
        max_wavenumber = max_wavenumber.to("1/m") * 2 * np.pi
        min_freq = min_freq.to("1/s") * 2 * np.pi
        max_freq = max_freq.to("1/s") * 2 * np.pi

        self.min_depth = min_depth
        self.max_depth = max_depth
        self.min_freq = min_freq
        self.max_freq = max_freq
        self.min_wavenumber = min_wavenumber
        self.max_wavenumber = max_wavenumber
        self.g = g
        self.beta = beta

        self.k_max_lim = Q_(50, 'zonal_wavenumber').to('1/m') * 2 * np.pi
        self.k_min_lim = Q_(-50, 'zonal_wavenumber').to('1/m') * 2 * np.pi
        self.omega_max_lim = Q_(2, '1/day').to('1/s') * 2 * np.pi
        self.omega_min_lim = Q_(-0.5, '1/day').to('1/s') * 2 * np.pi

        self.shape: Polygon = None

        self.border_shape = Polygon([(self.min_wavenumber.m, self.min_freq.m),
                                (self.max_wavenumber.m, self.min_freq.m),
                                (self.max_wavenumber.m, self.max_freq.m),
                                (self.min_wavenumber.m, self.max_freq.m)])
        self.desc = ""

    def shape_back_to_wavenumber_cpd(self):
        x = self.shape.exterior.xy[0]
        y = self.shape.exterior.xy[1]
        x = np.array(x)
        y = np.array(y)
        x = x * EARTH_EQUATOR_LENGTH.m / (2 * np.pi)
        y = y * 86400 / (2 * np.pi)
        self.shape = Polygon(np.array([x, y]).T)

class KelvinFilterZone(FilterZone):
    def __init__(self, min_depth: Quantity|float=8, max_depth: Quantity|float=90,
                 min_freq: Quantity|float=1/30, max_freq: Quantity|float=0.4,
                 min_wavenumber: Quantity|float=1,
                 max_wavenumber: Quantity|float=14,
                 g: Quantity|float=None, beta: Quantity|float=None):
        super().__init__(min_depth, max_depth, min_freq, max_freq,
                         min_wavenumber, max_wavenumber, g, beta)

        self.shape = Polygon([(0, 0),
                    (self.k_max_lim.m, np.sqrt(self.g * self.min_depth).m*self.k_max_lim.m),
                    (self.k_max_lim.m, np.sqrt(self.g * self.max_depth).m*self.k_max_lim.m),
        ])
        self.shape = self.shape.intersection(self.border_shape)
        self.shape_back_to_wavenumber_cpd()
    
        self.desc = json.dumps({
            "name": "Kelvin Filter Zone",
            "min_depth": self.min_depth.to('m').m,
            "max_depth": self.max_depth.to('m').m,
            "min_freq": self.min_freq.to('1/day').m / PI2,
            "max_freq": self.max_freq.to('1/day').m / PI2,
            "min_wavenumber": self.min_wavenumber.to('zonal_wavenumber').m / PI2,
            "max_wavenumber": self.max_wavenumber.to('zonal_wavenumber').m / PI2,
            "g": self.g.to('m/s^2').m,
            "beta": self.beta.to('1/m/s').m,
            "description": "Kelvin wave filter zone in the zonal-wavenumber frequency domain."
        })


class RossbyFilterZone(FilterZone):
    def __init__(self, n :int=1, min_depth: Quantity|float=8, max_depth: Quantity|float=90,
                 min_freq: Quantity|float=1/60, max_freq: Quantity|float=0.2,
                 min_wavenumber: Quantity|float=-10,
                 max_wavenumber: Quantity|float=0,
                 g: Quantity|float=None, beta: Quantity|float=None):
        super().__init__(min_depth, max_depth, min_freq, max_freq,
                         min_wavenumber, max_wavenumber, g, beta)

        k = np.linspace(self.max_wavenumber.m, self.min_wavenumber.m, 100)
        omega1 = self.beta.m * -k / (k**2+self.beta.m/np.sqrt(self.g.m*self.max_depth.m)*(2*n+1))
        omega2 = self.beta.m * -k / (k**2+self.beta.m/np.sqrt(self.g.m*self.min_depth.m)*(2*n+1))
        line1 = np.array([k, omega1]).T
        line2 = np.array([k[::-1], omega2[::-1]]).T
        self.shape = Polygon(np.concatenate([line1, line2], axis=0))
        self.shape = self.shape.intersection(self.border_shape)
        self.shape_back_to_wavenumber_cpd()

        self.desc = json.dumps({
            "name": "Equatorial Rossby Filter Zone",
            "n": n,
            "min_depth": self.min_depth.to('m').m,
            "max_depth": self.max_depth.to('m').m,
            "min_freq": self.min_freq.to('1/day').m / PI2,
            "max_freq": self.max_freq.to('1/day').m / PI2,
            "min_wavenumber": self.min_wavenumber.to('zonal_wavenumber').m / PI2,
            "max_wavenumber": self.max_wavenumber.to('zonal_wavenumber').m / PI2,
            "g": self.g.to('m/s^2').m,
            "beta": self.beta.to('1/m/s').m,
            "description": f"Equatorial Rossby wave filter zone in the"
            "zonal-wavenumber frequency domain with n={n}."
        })


class MRGFilterZone(FilterZone):
    def __init__(self, min_depth: Quantity|float=8, max_depth: Quantity|float=90,
                 min_freq: Quantity|float=0.1, max_freq: Quantity|float=0.35,
                 min_wavenumber: Quantity|float=-10,
                 max_wavenumber: Quantity|float=-1,
                 g: Quantity|float=None, beta: Quantity|float=None):
        super().__init__(min_depth, max_depth, min_freq, max_freq,
                         min_wavenumber, max_wavenumber, g, beta)

        k = np.linspace(self.max_wavenumber.m, self.min_wavenumber.m, 100)
        c1 = np.sqrt(self.g.m * self.max_depth.m)
        c2 = np.sqrt(self.g.m * self.min_depth.m)
        omega1 = np.sqrt(k**2*c1**2/4+c1*self.beta.m) - -k*c1/2
        omega2 = np.sqrt(k**2*c2**2/4+c2*self.beta.m) - -k*c2/2
        line1 = np.array([k, omega1]).T
        line2 = np.array([k[::-1], omega2[::-1]]).T
        self.shape = Polygon(np.concatenate([line1, line2], axis=0))
        self.shape = self.shape.intersection(self.border_shape)
        self.shape_back_to_wavenumber_cpd()

        self.desc = json.dumps({
            "name": "Mixed Rossby Gravity Filter Zone",
            "min_depth": self.min_depth.to('m').m,
            "max_depth": self.max_depth.to('m').m,
            "min_freq": self.min_freq.to('1/day').m / PI2,
            "max_freq": self.max_freq.to('1/day').m / PI2,
            "min_wavenumber": self.min_wavenumber.to('zonal_wavenumber').m / PI2,
            "max_wavenumber": self.max_wavenumber.to('zonal_wavenumber').m / PI2,
            "g": self.g.to('m/s^2').m,
            "beta": self.beta.to('1/m/s').m,
            "description": "Mixed Rossby Gravity wave filter zone in the"
            "zonal-wavenumber frequency domain."
        })


class EIGFilterZone(FilterZone):
    def __init__(self, n:int=0, min_depth: Quantity|float=12, max_depth: Quantity|float=50,
                 min_freq: Quantity|float=0.2, max_freq: Quantity|float=0.55,
                 min_wavenumber: Quantity|float=1,
                 max_wavenumber: Quantity|float=14,
                 g: Quantity|float=None, beta: Quantity|float=None):
        super().__init__(min_depth, max_depth, min_freq, max_freq,
                         min_wavenumber, max_wavenumber, g, beta)

        k = np.linspace(self.min_wavenumber.m, self.k_max_lim.m, 100)
        c1 = np.sqrt(self.g.m * self.max_depth.m)
        c2 = np.sqrt(self.g.m * self.min_depth.m)
        if n == 0:
            omega1 = np.sqrt(k**2*c1**2/4+c1*self.beta.m) - -k*c1/2
            omega2 = np.sqrt(k**2*c2**2/4+c2*self.beta.m) - -k*c2/2
            line1 = np.array([k, omega1]).T
            line2 = np.array([k[::-1], omega2[::-1]]).T
        elif n > 0:
            omega1 = np.sqrt(k**2*c1**2+c1*self.beta.m*(2*n+1))
            omega2 = np.sqrt(k**2*c2**2+c2*self.beta.m*(2*n+1))
            line1 = np.array([k, omega1]).T
            line2 = np.array([k[::-1], omega2[::-1]]).T
        else:
            raise ValueError("n should be greater than 0")
        self.shape = Polygon(np.concatenate([line1, line2], axis=0))
        self.shape = self.shape.intersection(self.border_shape)
        self.shape_back_to_wavenumber_cpd()

        self.desc = json.dumps({
            "name": "Equatorial Inertia Gravity Filter Zone",
            "n": n,
            "min_depth": self.min_depth.to('m').m,
            "max_depth": self.max_depth.to('m').m,
            "min_freq": self.min_freq.to('1/day').m / PI2,
            "max_freq": self.max_freq.to('1/day').m / PI2,
            "min_wavenumber": self.min_wavenumber.to('zonal_wavenumber').m / PI2,
            "max_wavenumber": self.max_wavenumber.to('zonal_wavenumber').m / PI2,
            "g": self.g.to('m/s^2').m,
            "beta": self.beta.to('1/m/s').m,
            "description": f"Equatorial Inertia Gravity wave filter zone in the"
            "zonal-wavenumber frequency domain with n={n}."
        })


class WIGFilterZone(FilterZone):
    def __init__(self, n:int=1, min_depth: Quantity|float=12, max_depth: Quantity|float=50,
                 min_freq: Quantity|float=0.3, max_freq: Quantity|float=0.7,
                 min_wavenumber: Quantity|float=-15,
                 max_wavenumber: Quantity|float=-1,
                 g: Quantity|float=None, beta: Quantity|float=None):
        if n < 1:
            raise ValueError("n should be greater or equal to 1")

        super().__init__(min_depth, max_depth, min_freq, max_freq,
                         min_wavenumber, max_wavenumber, g, beta)

        k = np.linspace(self.max_wavenumber.m, self.min_wavenumber.m, 100)
        c1 = np.sqrt(self.g.m * self.max_depth.m)
        c2 = np.sqrt(self.g.m * self.min_depth.m)
        omega1 = np.sqrt(k**2*c1**2+c1*self.beta.m*(2*n+1))
        omega2 = np.sqrt(k**2*c2**2+c2*self.beta.m*(2*n+1))
        line1 = np.array([k, omega1]).T
        line2 = np.array([k[::-1], omega2[::-1]]).T
        self.shape = Polygon(np.concatenate([line1, line2], axis=0))
        self.shape = self.shape.intersection(self.border_shape)
        self.shape_back_to_wavenumber_cpd()

        self.desc = json.dumps({
            "name": "Westward Inertia Gravity Filter Zone",
            "n": n,
            "min_depth": self.min_depth.to('m').m,
            "max_depth": self.max_depth.to('m').m,
            "min_freq": self.min_freq.to('1/day').m / PI2,
            "max_freq": self.max_freq.to('1/day').m / PI2,
            "min_wavenumber": self.min_wavenumber.to('zonal_wavenumber').m / PI2,
            "max_wavenumber": self.max_wavenumber.to('zonal_wavenumber').m / PI2,
            "g": self.g.to('m/s^2').m,
            "beta": self.beta.to('1/m/s').m,
            "description": f"Westward Inertia Gravity wave filter zone in the"
            "zonal-wavenumber frequency domain with n={n}."
        })


class MJOFilterZone(FilterZone):
    def __init__(self, min_freq: Quantity|float=1/96, 
                 max_freq: Quantity|float=1/30,
                 min_wavenumber: Quantity|float=0.5,
                 max_wavenumber: Quantity|float=5):

        super().__init__(0, 1, min_freq, max_freq,
                         min_wavenumber, max_wavenumber, 0, 0)

        coord = [[self.min_wavenumber.m, self.max_freq.m],
                 [self.max_wavenumber.m, self.max_freq.m],
                 [self.max_wavenumber.m, self.min_freq.m],
                 [self.min_wavenumber.m, self.min_freq.m]]
        self.shape = Polygon(coord)
        self.shape = self.shape.intersection(self.border_shape)
        self.shape_back_to_wavenumber_cpd()
        
        self.desc = json.dumps({
            "name": "MJO Filter Zone",
            "min_freq": self.min_freq.to('1/day').m / PI2,
            "max_freq": self.max_freq.to('1/day').m / PI2,
            "min_wavenumber": self.min_wavenumber.to('zonal_wavenumber').m / PI2,
            "max_wavenumber": self.max_wavenumber.to('zonal_wavenumber').m / PI2,
            "description": "MJO wave filter zone in the zonal-wavenumber frequency domain."
        })


class TDFilterZone(FilterZone):
    def __init__(self, min_freq: Quantity|float=1/6, 
                 max_freq: Quantity|float=1/2.5,
                 min_wavenumber: Quantity|float=-18,
                 max_wavenumber: Quantity|float=-8):

        super().__init__(0, 1, min_freq, max_freq,
                         min_wavenumber, max_wavenumber, 0, 0)

        coord = [[self.min_wavenumber.m, self.max_freq.m],
                 [self.max_wavenumber.m, self.max_freq.m],
                 [self.max_wavenumber.m, self.min_freq.m],
                 [self.min_wavenumber.m, self.min_freq.m]]
        self.shape = Polygon(coord)
        self.shape = self.shape.intersection(self.border_shape)
        self.shape_back_to_wavenumber_cpd()

        self.desc = json.dumps({
            "name": "Tropical Depression Filter Zone",
            "min_freq": self.min_freq.to('1/day').m / PI2,
            "max_freq": self.max_freq.to('1/day').m / PI2,
            "min_wavenumber": self.min_wavenumber.to('zonal_wavenumber').m / PI2,
            "max_wavenumber": self.max_wavenumber.to('zonal_wavenumber').m / PI2,
            "description": "Tropical Depression wave filter zone in the zonal-wavenumber frequency domain."
        })

def customFilterZone(wavenumber: Quantity | np.ndarray,
                    frequency: Quantity | np.ndarray,) -> Polygon:
    wavenumber = _handle_quantity(wavenumber, "zonal_wavenumber")
    frequency = _handle_quantity(frequency, "1/day")
    return Polygon(np.array([wavenumber.m, frequency.m]).T)


def fft2_zone_filter(data: xr.DataArray | np.ndarray,
                    filter_zone: FilterZone | Polygon,
                    axes: tuple[int] = (0, 1),
                    x_shift_freq: np.ndarray = None,
                    y_shift_freq: np.ndarray = None,):
    if isinstance(filter_zone, FilterZone):
        filter_zone = filter_zone.shape
    elif not isinstance(filter_zone, Polygon):
        raise TypeError("filter_zone must be a Polygon or FilterZone")

    if x_shift_freq is None:
        x_shift_freq = np.fft.fftfreq(data.shape[axes[0]])
        x_shift_freq = np.fft.fftshift(x_shift_freq)
    if y_shift_freq is None:
        y_shift_freq = np.fft.fftfreq(data.shape[axes[1]])
        y_shift_freq = np.fft.fftshift(y_shift_freq)

    data_fft = np.fft.fft2(data, axes=axes)
    data_fft = np.fft.fftshift(data_fft, axes=axes)

    # 创建mask
    X, Y = np.meshgrid(-x_shift_freq, y_shift_freq)
    points = spoints(X.flatten(), Y.flatten())
    mask1 = filter_zone.contains(points)
    mask1 = mask1.reshape(X.shape)
    # 对称部分
    x = np.array(filter_zone.exterior.xy[0])
    y = np.array(filter_zone.exterior.xy[1])
    filter_zone = Polygon(np.array([-x, -y]).T)
    mask2 = filter_zone.contains(points)
    mask2 = mask2.reshape(X.shape)
    mask = mask1 | mask2
    # 变为 false=0, true=1 的mask
    mask = mask.astype(np.float32)

    # 在axes的维度上乘上mask
    # 更换数据的维度
    mask_shape = [1] * len(data.shape)
    for i, axis in enumerate(axes):
        mask_shape[axis] = data.shape[axis]
    mask = mask.reshape(mask_shape)
    data_fft = data_fft * mask

    # 还原数据
    data_fft = np.fft.ifftshift(data_fft, axes=axes)
    data_fft = np.fft.ifft2(data_fft, axes=axes)

    return np.real(data_fft)


def wk99_filter(data: xr.DataArray, filter_zone: FilterZone | Polygon,
                method: Literal['none', 'sym', 'asym'] = 'none',
                lon_name: str = "longitude", time_name: str = "time",
                lat_name: str = "latitude"):
    from ..calc.xr import safe_sel
    dt = data[time_name].values[1] - data[time_name].values[0]
    dt = np.timedelta64(dt, 's')
    ticks_in_day = np.timedelta64(1, 'D') / dt

    lon = data[lon_name].values
    zonal_res = lon[1] - lon[0]

    lon_axis = data.get_axis_num(lon_name)
    time_axis = data.get_axis_num(time_name)

    if method in ['sym', 'asym']:
        data_north = safe_sel(data, {lat_name: slice(90, 0)})
        data_south = safe_sel(data, {lat_name: slice(0, -90)})
        lat = data_north[lat_name].values
        data_north = data_north.values
        data_south = data_south.values

        data_south = np.flip(data_south, axis=data.get_axis_num(lat_name))

        if method == 'sym':
            _data = (data_north + data_south) / 2
        else:
            _data = (data_north - data_south) / 2
    else:
        lat = data[lat_name].values
        _data = data.values

    x_shift_freq = np.fft.fftfreq(data.shape[lon_axis], np.abs(zonal_res) / 360)
    x_shift_freq = np.fft.fftshift(x_shift_freq)
    y_shift_freq = np.fft.fftfreq(data.shape[time_axis], 1/ticks_in_day)
    y_shift_freq = np.fft.fftshift(y_shift_freq)

    result =  fft2_zone_filter(_data, filter_zone, axes=(time_axis, lon_axis),
                            x_shift_freq=x_shift_freq, y_shift_freq=y_shift_freq)

    return xr.DataArray(result, dims=data.dims,
                            coords={lon_name: lon,
                                    time_name: data[time_name].values,
                                    lat_name: lat},
                            attrs=data.attrs)


def fft2_filter(data: xr.DataArray, filter_zone: FilterZone | Polygon,
                lon_name: str = "longitude", time_name: str = "time",
                lat_name: str = "latitude", time_window: bool=True,
                lon_window: bool=True):
    from scipy.signal.windows import tukey
    dt = data[time_name].values[1] - data[time_name].values[0]
    dt = np.timedelta64(dt, 's')
    ticks_in_day = np.timedelta64(1, 'D') / dt

    lon = data[lon_name].values
    zonal_res = lon[1] - lon[0]

    lon_axis = data.get_axis_num(lon_name)
    time_axis = data.get_axis_num(time_name)

    lat = data[lat_name].values
    _data = data.values

    if time_window:
        slices = [slice(None) if i == time_axis else None for i in range(_data.ndim)]
        _data = _data * tukey(_data.shape[time_axis], 0.1)[*slices]
    if lon_window:
        slices = [slice(None) if i == lon_axis else None for i in range(_data.ndim)]
        _data = _data * tukey(_data.shape[lon_axis], 0.1)[*slices]

    x_shift_freq = np.fft.fftfreq(data.shape[lon_axis], np.abs(zonal_res) / 360)
    x_shift_freq = np.fft.fftshift(x_shift_freq)
    y_shift_freq = np.fft.fftfreq(data.shape[time_axis], 1/ticks_in_day)
    y_shift_freq = np.fft.fftshift(y_shift_freq)

    result =  fft2_zone_filter(_data, filter_zone, axes=(time_axis, lon_axis),
                            x_shift_freq=x_shift_freq, y_shift_freq=y_shift_freq)

    return xr.DataArray(result, dims=data.dims,
                            coords={lon_name: lon,
                                    time_name: data[time_name].values,
                                    lat_name: lat},
                            attrs=data.attrs)